package com.ecmaster.search.count;

import org.apache.hadoop.mapreduce.RecordWriter;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.apache.hadoop.mapreduce.lib.output.FileOutputFormat;
import redis.clients.jedis.Jedis;

import java.io.IOException;

/**
 * Created by GanJc on 2015/9/15.
 */
public class RedisOutputFormat<K, V> extends FileOutputFormat<K, V> {

    protected static class RedisRecordWriter extends RecordWriter {

        private Jedis jedis  ;

        public RedisRecordWriter(Jedis jedis) {
            this.jedis = jedis;
        }

        @Override
        public void write(Object key, Object val) throws IOException, InterruptedException {
            if (key == null || val == null) {
                return;
            }
            String s = key.toString();
            for (int i = 0; i < s.length(); i++) {
                String k = s.substring(0, i + 1);
                int score = Integer.parseInt(val.toString());
                jedis.zincrby(k, score, s);
            }
        }

        @Override
        public void close(TaskAttemptContext taskAttemptContext) throws IOException, InterruptedException {
            if (jedis != null) {
                jedis.disconnect();
            }
        }
    }

    @Override
    public RecordWriter<K, V> getRecordWriter(TaskAttemptContext taskAttemptContext) throws IOException, InterruptedException {
        return new RedisRecordWriter(new Jedis("127.0.0.1"));
    }

}
